package com.xinmao.quantitative.utils;

import cn.hutool.core.date.DateUtil;
import org.jfree.chart.ChartPanel;
import org.jfree.chart.JFreeChart;
import org.jfree.chart.axis.NumberAxis;
import org.jfree.chart.plot.Marker;
import org.jfree.chart.plot.ValueMarker;
import org.jfree.chart.plot.XYPlot;
import org.jfree.chart.renderer.xy.StandardXYItemRenderer;
import org.jfree.chart.ui.ApplicationFrame;
import org.jfree.chart.ui.UIUtils;
import org.jfree.data.time.Day;
import org.jfree.data.time.Minute;
import org.jfree.data.time.TimeSeries;
import org.jfree.data.time.TimeSeriesCollection;
import org.ta4j.core.*;
import org.ta4j.core.num.Num;

import java.awt.*;
import java.util.Calendar;
import java.util.Date;

/**
 * @Title: ChartUtil
 * @Author: XYe13
 * @Date: 5/20/2024 10:27
 * @Description: <TODO description class purpose>
 */
public class ChartUtil {

    public static TimeSeries buildChartBarSeriesDay(BarSeries barSeries, Indicator<Num> indicator,
                                                                      String name) {
        TimeSeries chartBarSeries = new TimeSeries(name);
        for (int i = 0; i < barSeries.getBarCount(); i++) {
            Bar bar = barSeries.getBar(i);
            chartBarSeries.add(new Day(DateUtil.date(bar.getEndTime().toLocalDate())), indicator.getValue(i).doubleValue());
        }
        return chartBarSeries;
    }

    public static TimeSeries buildChartBarSeriesMinute(BarSeries barSeries, Indicator<Num> indicator, String name) {
        TimeSeries chartBarSeries = new TimeSeries(name);
        for (int i = 0; i < barSeries.getBarCount(); i++) {
            Bar bar = barSeries.getBar(i);
            chartBarSeries.add(new Minute(DateUtil.date(bar.getEndTime().toLocalDateTime()),Calendar.getInstance()),
                    indicator.getValue(i).doubleValue());
        }
        return chartBarSeries;
    }

    public static void addDataFlowAxis(XYPlot plot, TimeSeriesCollection dataset) {
        final NumberAxis cashAxis = new NumberAxis("Data Ratio");
        cashAxis.setAutoRangeIncludesZero(false);
        plot.setRangeAxis(1, cashAxis);
        plot.setDataset(1, dataset);
        plot.mapDatasetToRangeAxis(1, 1);
        final StandardXYItemRenderer cashFlowRenderer = new StandardXYItemRenderer();
        cashFlowRenderer.setSeriesPaint(0, Color.blue);
        plot.setRenderer(1, cashFlowRenderer);
    }


    public static void displayChart(JFreeChart chart) {
        // Chart panel
        ChartPanel panel = new ChartPanel(chart);
        panel.setFillZoomRectangle(true);
        panel.setMouseWheelEnabled(true);
        panel.setPreferredSize(new Dimension(1960, 600));
        // Application frame
        ApplicationFrame frame = new ApplicationFrame("Data chart");
        frame.setContentPane(panel);
        frame.pack();
        UIUtils.centerFrameOnScreen(frame);
        frame.setVisible(true);
    }

    /**
     * Runs a strategy over a bar series and adds the value markers corresponding to
     * buy/sell signals to the plot.
     *
     * @param series   the bar series
     * @param plot     the plot
     */
    public static void addBuySellSignals(BarSeries series,TradingRecord tradingRecord, XYPlot plot) {
        // Adding markers to plot
        for (Position position : tradingRecord.getPositions()) {
            // Buy signal
            double buySignalBarTime = new Minute(
                    Date.from(series.getBar(position.getEntry().getIndex()).getEndTime().toInstant()))
                    .getFirstMillisecond();
            Marker buyMarker = new ValueMarker(buySignalBarTime);
            buyMarker.setPaint(Color.GREEN);
            buyMarker.setLabel("B");
            plot.addDomainMarker(buyMarker);
            // Sell signal
            double sellSignalBarTime = new Minute(
                    Date.from(series.getBar(position.getExit().getIndex()).getEndTime().toInstant()))
                    .getFirstMillisecond();
            Marker sellMarker = new ValueMarker(sellSignalBarTime);
            sellMarker.setPaint(Color.RED);
            sellMarker.setLabel("S");
            plot.addDomainMarker(sellMarker);
        }
    }
}
